package com.ppch.wuwamanus.advisors;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

/**
 * Package:ppch-ai-agent
 * ClassName:MyLoggerAdvisor
 *
 * @Author 泡泡茶壶
 * @Create 2025/5/6 15:15
 * @Version 1.0
 * Description:
 * 自定义日志记录顾问（拦截器）
 */
@Slf4j
public class MyLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    /**
     * 提供唯一的 advisor 名称。
     */
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    /**
     * 确定链中的 advisor 顺序（值越小优先级越高）
     */
    @Override
    public int getOrder() {
        return 0;
    }

    /**
     * 前置顾问
     */
    private AdvisedRequest before(AdvisedRequest request) {
        //log.info("MyLoggerAdvisor打印日志，用户输入: {}", request.userText());
        return request;
    }

    /**
     * 后置顾问
     */
    private void observeAfter(AdvisedResponse advisedResponse) {
        //获取响应文本
        String responseText = advisedResponse.response().getResult().getOutput().getText();
        Integer totalTokens = advisedResponse.response().getMetadata().getUsage().getTotalTokens();
        Integer promptTokens = advisedResponse.response().getMetadata().getUsage().getPromptTokens();
        Integer completionTokens = totalTokens-promptTokens;
        log.info("MyLoggerAdvisor打印日志，消耗总Tokens：{}, 用户提示消耗Tokens：{}, AI回复消耗Tokens：{}",
                totalTokens,promptTokens,completionTokens);
        log.info("MyLoggerAdvisor打印日志，AI响应: {}", responseText);
    }

    /**
     * 针对非流式处理
     * @param advisedRequest 请求
     * @param chain advisors 链
     * @return
     */
    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        //1.前置顾问
        advisedRequest = this.before(advisedRequest);
        //2.调用链中的下一个advisor
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        //3.后置顾问
        this.observeAfter(advisedResponse);
        return advisedResponse;
    }

    /**
     * 针对流式处理
     * @param advisedRequest 请求
     * @param chain advisors 链
     * @return
     */
    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        //1.前置顾问
        advisedRequest = this.before(advisedRequest);
        //2.调用链中的下一个advisor
        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);
        //3.将 Flux 响应聚合到单个 AdvisedResponse 中
        MessageAggregator messageAggregator = new MessageAggregator();
        Flux<AdvisedResponse> advisedResponseFlux = messageAggregator.aggregateAdvisedResponse(advisedResponses, this::observeAfter);
        return advisedResponseFlux;
    }
}
